package com.DL4J;

import org.deeplearning4j.datasets.iterator.impl.ListDataSetIterator;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.Sgd;
import org.nd4j.linalg.lossfunctions.LossFunctions;
import org.apache.poi.ss.usermodel.*;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;


import java.io.*;
import java.util.Collections;
import java.util.List;

public class MultivariateLinearRegressionDL4J {

    private static final String MODEL_FILE_PATH = "D:\\xibike\\回归测试\\model\\modelCoefficients451.zip";
    private static final String EXCEL_FILE_PATH = "D:\\xibike\\回归测试\\回归算法测试数据.xlsx";

    public static void main(String[] args) throws Exception {
        // 如果模型文件不存在，则训练并保存模型
        if (!new File(MODEL_FILE_PATH).exists()) {
            trainAndSaveModel();
        }

        // 加载模型
        File modelFile = new File(MODEL_FILE_PATH);
        MultiLayerNetwork model = MultiLayerNetwork.load(modelFile, true);

        // 预测新数据点 回水 供水 室外温度 开度
        double[] newX = {35.46, 39.79, 9, 100};
        // 归一化处理
        newX[3] = normalize(newX[3], 0, 100);
        // 创建形状为 [1, 5] 的输入矩阵
        INDArray input = Nd4j.create(new double[][]{{1.0, newX[0], newX[1], newX[2], newX[3]}}); // 使用二维数组
        INDArray output = model.output(input);
        System.out.println("预测值: " + output.getDouble(0));
    }

    private static void trainAndSaveModel() throws IOException {
        // 读取 Excel 文件中的特征值和目标变量
        DataSetIterator dataIterator = readExcelFile(EXCEL_FILE_PATH);

        // 构建模型配置
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
                .updater(new Sgd(0.01))
                .list()
                .layer(0, new DenseLayer.Builder().nIn(5).nOut(10)
                        .activation(Activation.RELU).build())
                .layer(1, new OutputLayer.Builder(LossFunctions.LossFunction.MSE)
                        .activation(Activation.IDENTITY).nIn(10).nOut(1).build())
                .build();

        MultiLayerNetwork model = new MultiLayerNetwork(conf);
        model.init();
        model.setListeners(new ScoreIterationListener(10));

        // 训练模型
        int numEpochs = 1000; // 设置迭代次数
        for (int i = 0; i < numEpochs; i++) {
            while (dataIterator.hasNext()) {
                DataSet next = dataIterator.next();
                model.fit(next);
            }
            dataIterator.reset();
        }

        // 保存模型
        model.save(new File(MODEL_FILE_PATH), true);
    }

    private static DataSetIterator readExcelFile(String filePath) throws IOException {
        try (FileInputStream fis = new FileInputStream(filePath);
             Workbook workbook = new XSSFWorkbook(fis)) {

            Sheet sheet = workbook.getSheetAt(0);
            int rowCount = sheet.getLastRowNum() + 1;
            int columnCount = 5; // 包括截距项

            INDArray features = Nd4j.zeros(rowCount - 1, columnCount); // 减去表头行
            INDArray labels = Nd4j.zeros(rowCount - 1, 1);

            for (int i = 1; i < rowCount; i++) {
                Row row = sheet.getRow(i);
                for (int j = 0; j < columnCount - 1; j++) {
                    Cell cell = row.getCell(j);
                    features.putScalar(new int[]{i - 1, j + 1}, cell == null ? 0.0 : cell.getNumericCellValue());
                }
                Cell targetCell = row.getCell(columnCount - 1);
                labels.putScalar(new int[]{i - 1, 0}, targetCell == null ? 0.0 : targetCell.getNumericCellValue());

                // 对阀门开度进行归一化处理
                features.putScalar(new int[]{i - 1, columnCount - 1}, normalize(features.getDouble(i - 1, columnCount - 1), 0, 100));
            }

            // 添加截距项（即常数列）
            INDArray intercept = Nd4j.ones(rowCount - 1, 1);
            features.putColumn(0, intercept); // 将截距项放在第一列

            // 创建一个包含单个DataSet的列表
            List<DataSet> dataSets = Collections.singletonList(new DataSet(features, labels));

            // 返回一个基于该列表的DataSetIterator
            return new ListDataSetIterator(dataSets);
        }
    }

    private static double normalize(double value, double min, double max) {
        return (value - min) / (max - min);
    }
}